{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "761319ec-18e8-46c5-b7d6-bfdabc892e38",
   "metadata": {},
   "source": [
    "## Post-process a finetuned LLM\n",
    "\n",
    "Test and upload a finetuned language model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "00e6173c-597b-4eb0-8275-bab2c62d4d48",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -q -U huggingface_hub peft transformers torch accelerate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d9e47e4-3adf-4aa4-910c-9f75eb213420",
   "metadata": {},
   "outputs": [],
   "source": [
    "from huggingface_hub import notebook_login\n",
    "\n",
    "notebook_login()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25376737-3af9-4382-9240-9e94050c8ae7",
   "metadata": {},
   "outputs": [],
   "source": [
    "!git config --global credential.helper store"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e24d7ddd-044f-4739-8b01-5a1b187150f5",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63449bd1-4434-4f0e-b41c-e3d86815319b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from peft import PeftModel, PeftConfig\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bfeb06bd-96d9-4db0-81ee-1053104f915c",
   "metadata": {},
   "outputs": [],
   "source": [
    "!nvidia-smi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2057759-7bd5-460a-86fd-c49897fbf79c",
   "metadata": {},
   "outputs": [],
   "source": [
    "free_in_GB = int(torch.cuda.mem_get_info()[0] / 1024**3)\n",
    "max_memory = f\"{free_in_GB-2}GB\"\n",
    "n_gpus = torch.cuda.device_count()\n",
    "max_memory = {i: max_memory for i in range(n_gpus)}\n",
    "print(\"max memory: \", max_memory)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c68a3bba-f374-4ed6-acb6-ebedeb6e29be",
   "metadata": {},
   "source": [
    "## Loss curve\n",
    "\n",
    "During training, the model converged nicely as follows:\n",
    "\n",
    "![image](https://raw.githubusercontent.com/daniel-furman/sft-demos/main/assets/jul_24_23_1_14_00_log_loss_curves_llama-2-70b-dolphin.png)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d6076740-c324-49a9-841e-30b4878d208b",
   "metadata": {},
   "source": [
    "## Basic testing\n",
    "\n",
    "With a supervised finetuned (sft) model in hand, we can test it on some basic prompts and then upload it to the Hugging Face hub either as a public or private model repo, depending on the use case."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b17e2d0-b482-4ebd-a830-325d98318f70",
   "metadata": {},
   "outputs": [],
   "source": [
    "peft_model_id = \"results/checkpoint-5000\"\n",
    "# peft_model_id = \"dfurman/llama-2-70b-dolphin-peft\"\n",
    "config = PeftConfig.from_pretrained(peft_model_id)\n",
    "\n",
    "bnb_config = BitsAndBytesConfig(\n",
    "    load_in_4bit=True,\n",
    "    bnb_4bit_quant_type=\"nf4\",\n",
    "    bnb_4bit_compute_dtype=torch.bfloat16,\n",
    ")\n",
    "\n",
    "model = AutoModelForCausalLM.from_pretrained(\n",
    "    config.base_model_name_or_path,\n",
    "    quantization_config=bnb_config,\n",
    "    use_auth_token=True,\n",
    "    torch_dtype=torch.bfloat16,\n",
    "    device_map=\"auto\",\n",
    ")\n",
    "tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)\n",
    "tokenizer.pad_token = tokenizer.eos_token\n",
    "\n",
    "# Load the Lora model\n",
    "model = PeftModel.from_pretrained(model, peft_model_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b93b4abc-50e6-4308-8bfb-dfd164aef43f",
   "metadata": {},
   "outputs": [],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "136794be-1cf3-4478-90df-61507179f076",
   "metadata": {},
   "outputs": [],
   "source": [
    "free_in_GB = int(torch.cuda.mem_get_info()[0] / 1024**3)\n",
    "max_memory = f\"{free_in_GB-2}GB\"\n",
    "n_gpus = torch.cuda.device_count()\n",
    "max_memory = {i: max_memory for i in range(n_gpus)}\n",
    "print(\"max memory: \", max_memory)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "175e9207-f17b-4ba0-9d3c-cf1ad2f69946",
   "metadata": {},
   "outputs": [],
   "source": [
    "# text generation function\n",
    "\n",
    "\n",
    "def falcon_generate(\n",
    "    model: AutoModelForCausalLM,\n",
    "    tokenizer: AutoTokenizer,\n",
    "    prompt: str,\n",
    "    max_new_tokens: int = 128,\n",
    "    temperature: int = 1.0,\n",
    ") -> str:\n",
    "    \"\"\"\n",
    "    Initialize the pipeline\n",
    "    Uses Hugging Face GenerationConfig defaults\n",
    "        https://huggingface.co/docs/transformers/v4.29.1/en/main_classes/text_generation#transformers.GenerationConfig\n",
    "    Args:\n",
    "        model (transformers.AutoModelForCausalLM): Falcon model for text generation\n",
    "        tokenizer (transformers.AutoTokenizer): Tokenizer for model\n",
    "        prompt (str): Prompt for text generation\n",
    "        max_new_tokens (int, optional): Max new tokens after the prompt to generate. Defaults to 128.\n",
    "        temperature (float, optional): The value used to modulate the next token probabilities.\n",
    "            Defaults to 1.0\n",
    "    \"\"\"\n",
    "    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "    inputs = tokenizer(\n",
    "        [prompt],\n",
    "        return_tensors=\"pt\",\n",
    "        return_token_type_ids=False,\n",
    "    ).to(\n",
    "        device\n",
    "    )  # tokenize inputs, load on device\n",
    "\n",
    "    # when running Torch modules in lower precision, it is best practice to use the torch.autocast context manager.\n",
    "    with torch.autocast(\"cuda\", dtype=torch.bfloat16):\n",
    "        response = model.generate(\n",
    "            **inputs,\n",
    "            max_new_tokens=max_new_tokens,\n",
    "            temperature=temperature,\n",
    "            return_dict_in_generate=True,\n",
    "            eos_token_id=tokenizer.eos_token_id,\n",
    "            pad_token_id=tokenizer.pad_token_id,\n",
    "        )\n",
    "\n",
    "    decoded_output = tokenizer.decode(\n",
    "        response[\"sequences\"][0],\n",
    "        skip_special_tokens=True,\n",
    "    )  # grab output in natural language\n",
    "\n",
    "    return decoded_output[len(prompt) :]  # remove prompt from output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd8d36ba-e1a3-444b-b718-af104e130858",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = \"You are a helpful assistant. Write me a numbered list of things to do in New York City.\\n\"\n",
    "\n",
    "response = falcon_generate(\n",
    "    model,\n",
    "    tokenizer,\n",
    "    prompt,\n",
    "    max_new_tokens=250,\n",
    "    temperature=0.92,\n",
    ")\n",
    "\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b9753d7f-a9f9-4e96-ac7d-28cd1ef3034b",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = \"You are a helpful assistant. Daniel is in need of a haircut. His barber works Mondays, Wednesdays, and Fridays. So, Daniel went in for a haircut on Sunday. Does this make logical sense? Respond only with a yes or no answer in as few words as possible.\\n\"\n",
    "\n",
    "response = falcon_generate(\n",
    "    model,\n",
    "    tokenizer,\n",
    "    prompt,\n",
    "    max_new_tokens=100,\n",
    "    temperature=0.92,\n",
    ")\n",
    "\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3c7f8e8-aa1e-4cec-a6ca-20776da1673d",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = \"You are a helpful assistant. Write a short email inviting my friends to a dinner party on Friday. Respond succinctly.\\n\"\n",
    "\n",
    "response = falcon_generate(\n",
    "    model,\n",
    "    tokenizer,\n",
    "    prompt,\n",
    "    max_new_tokens=200,\n",
    "    temperature=0.92,\n",
    ")\n",
    "\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3f802dad-2bcf-4720-a30c-b3e7f89680ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = \"You are a helpful assistant. Tell me a recipe for vegan banana bread.\\n\"\n",
    "\n",
    "response = falcon_generate(\n",
    "    model,\n",
    "    tokenizer,\n",
    "    prompt,\n",
    "    max_new_tokens=500,\n",
    "    temperature=0.92,\n",
    ")\n",
    "\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f3ce0ab-d32e-4a7f-91e1-7bd88efc532b",
   "metadata": {},
   "source": [
    "## Inf runtime test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d0d1d4a3-e8d3-41f5-92eb-38b1b00eec29",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tqdm\n",
    "import time\n",
    "\n",
    "prompt = \"You are a helpful assistant. Write me a long list of things to do in San Francisco:\\n\"\n",
    "\n",
    "runtimes = []\n",
    "for i in tqdm.tqdm(range(25)):\n",
    "    start = time.time()\n",
    "    response = falcon_generate(\n",
    "        model,\n",
    "        tokenizer,\n",
    "        prompt,\n",
    "        max_new_tokens=50,\n",
    "        temperature=0.92,\n",
    "    )\n",
    "    end = time.time()\n",
    "    runtimes.append(end - start)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c381a747-4b96-4249-b6ca-1ac1ca7704d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "avg_runtime = torch.mean(torch.tensor(runtimes)).item()\n",
    "print(f\"Runtime avg in seconds: {avg_runtime}\")  # time in seconds"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "96ee6bc1-334f-4bc9-acdf-040ad9a532fd",
   "metadata": {},
   "source": [
    "## Upload model to Hugging Face\n",
    "1. Before running the cells below, create a model on your Hugging Face account. It can be a private or public repo and work with the below code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1972543-4563-47bf-9950-7dc8551b00db",
   "metadata": {},
   "outputs": [],
   "source": [
    "# push to hub\n",
    "model_id_load = \"dfurman/falcon-180b-instruct-peft\"\n",
    "\n",
    "# tokenizer\n",
    "tokenizer.push_to_hub(model_id_load, use_auth_token=True)\n",
    "# safetensors\n",
    "model.push_to_hub(model_id_load, use_auth_token=True, safe_serialization=True)\n",
    "# torch tensors\n",
    "model.push_to_hub(model_id_load, use_auth_token=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ba39e00",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "301faebbd5cea7fd4466786a19f1bea9d8baf657aaca95ef39840c46b8697603"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
